import os
import shutil
from pathlib import Path


def organize_cyclegan_images(source_dir, target_base_dir):
    """
    从CycleGAN/pix2pix结果文件夹中整理fake_B和real_B图片

    参数:
        source_dir: 包含所有图片的源文件夹路径
        target_base_dir: 目标基础文件夹路径
    """

    # 定义目标文件夹路径
    fake_b_dir = Path(target_base_dir) / "fake_images_diffusion"
    real_b_dir = Path(target_base_dir) / "real_images_diffusion"

    # 创建目标文件夹（如果不存在）
    fake_b_dir.mkdir(parents=True, exist_ok=True)
    real_b_dir.mkdir(parents=True, exist_ok=True)

    # 支持的图片格式
    image_extensions = {'.png', '.jpg', '.jpeg', '.tif', '.tiff', '.bmp'}

    # 计数器
    fake_b_count = 0
    real_b_count = 0
    other_files = 0

    print(f"开始从 {source_dir} 整理图片...")
    print("-" * 50)

    # 遍历源文件夹中的所有文件
    for filename in os.listdir(source_dir):
        file_path = os.path.join(source_dir, filename)

        # 只处理图片文件
        if os.path.isfile(file_path) and Path(filename).suffix.lower() in image_extensions:

            # 根据文件名模式分类
            if 'fake_B' in filename:
                # 复制fake_B图片到fake_images文件夹
                target_path = fake_b_dir / filename
                shutil.copy2(file_path, target_path)
                fake_b_count += 1
                print(f"✓ 复制 fake_B: {filename}")

            elif 'real_B' in filename:
                # 复制real_B图片到real_images文件夹
                target_path = real_b_dir / filename
                shutil.copy2(file_path, target_path)
                real_b_count += 1
                print(f"✓ 复制 real_B: {filename}")

            else:
                other_files += 1
                print(f"○ 跳过其他文件: {filename}")

    print("-" * 50)
    print("整理完成！")
    print(f"fake_B 图片数量: {fake_b_count}")
    print(f"real_B 图片数量: {real_b_count}")
    print(f"其他文件数量: {other_files}")
    print(f"fake_B 图片保存在: {fake_b_dir}")
    print(f"real_B 图片保存在: {real_b_dir}")


def main():
    # 配置路径 - 请根据您的实际情况修改这些路径
    source_directory = r"I:\lxj\stainedHE\pytorch_CycleGAN_and_pix2pix_master\results\diffusionwithvit\traindemo05_diffusionwithvit\test_600\images"
    target_directory = r"I:\lxj\stainedHE\pytorch_CycleGAN_and_pix2pix_master\virtual_staining_eval\data"

    # 检查源文件夹是否存在
    if not os.path.exists(source_directory):
        print(f"错误: 源文件夹不存在 - {source_directory}")
        return

    # 执行整理操作
    organize_cyclegan_images(source_directory, target_directory)


if __name__ == "__main__":
    main()